import cv2
import numpy as np
import torch
import time
import threading
import logging
from datetime import datetime
import os
from collections import deque

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# ---------------------- 1. 传感单元模拟 ----------------------
class CameraSensor:
    """模拟摄像头传感单元"""
    def __init__(self, video_source=0, fps=30):
        self.video_source = video_source
        self.fps = fps
        self.cap = None
        self.frame_width = 0
        self.frame_height = 0
        
    def open(self):
        """打开摄像头或视频文件"""
        self.cap = cv2.VideoCapture(self.video_source)
        if not self.cap.isOpened():
            raise ValueError(f"无法打开视频源: {self.video_source}")
            
        self.frame_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.frame_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        logger.info(f"摄像头已打开: {self.frame_width}x{self.frame_height} @ {self.fps}fps")
        
    def read_frame(self):
        """读取一帧图像"""
        ret, frame = self.cap.read()
        if not ret:
            logger.warning("无法读取帧，可能已到达视频末尾")
            return None
        return frame
    
    def release(self):
        """释放摄像头资源"""
        if self.cap:
            self.cap.release()
            logger.info("摄像头资源已释放")


# ---------------------- 2. 目标检测模型 ----------------------
class YOLOv5Detector:
    """YOLOv5目标检测模型"""
    def __init__(self, model_path='yolov5s.pt', device='cuda'):
        self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path)
        self.model.to(device)
        self.device = device
        self.classes = self.model.names
        
    def detect(self, frame, conf_threshold=0.5):
        """检测图像中的目标"""
        results = self.model(frame)
        detections = results.pandas().xyxy[0]
        detections = detections[detections['confidence'] > conf_threshold]
        
        boxes = detections[['xmin', 'ymin', 'xmax', 'ymax']].values.astype(int)
        scores = detections['confidence'].values
        class_ids = detections['class'].values
        class_names = [self.classes[int(cls)] for cls in class_ids]
        
        return boxes, scores, class_ids, class_names


# ---------------------- 3. 行为识别模型 ----------------------
class ActionRecognizer:
    """基于3D CNN的行为识别模型"""
    def __init__(self, model_path='i3d_r50.pth', device='cuda'):
        self.model = torch.load(model_path)
        self.model.to(device)
        self.model.eval()
        self.device = device
        self.buffer_size = 16  # 连续16帧作为输入
        self.transform = torch.nn.Sequential(
            torchvision.transforms.Resize((112, 112)),
            torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        )
        
    def recognize(self, frame_buffer):
        """识别行为"""
        if len(frame_buffer) < self.buffer_size:
            return 0  # 正常状态
        
        # 准备输入序列
        input_frames = []
        for frame in frame_buffer[-self.buffer_size:]:
            frame = cv2.resize(frame, (112, 112))
            frame = torch.from_numpy(frame).permute(2, 0, 1).float() / 255.0
            frame = self.transform(frame)
            input_frames.append(frame)
            
        input_tensor = torch.stack(input_frames).unsqueeze(0).to(self.device)  # [1, T, C, H, W]
        
        # 模型推理
        with torch.no_grad():
            output = self.model(input_tensor)
            _, predicted = torch.max(output, 1)
            
        return predicted.item()  # 0:正常, 1:打架, 2:摔倒, 3:奔跑


# ---------------------- 4. 烟火检测模型 ----------------------
class FireSmokeDetector:
    """烟火检测模型"""
    def __init__(self, model_path='fire_smoke_model.pt', device='cuda'):
        self.model = torch.load(model_path)
        self.model.to(device)
        self.model.eval()
        self.device = device
        
    def detect(self, frame, threshold=0.8):
        """检测烟火"""
        frame = cv2.resize(frame, (224, 224))
        frame = torch.from_numpy(frame).permute(2, 0, 1).float() / 255.0
        frame = frame.unsqueeze(0).to(self.device)
        
        with torch.no_grad():
            output = self.model(frame)
            prob = torch.sigmoid(output).item()
            
        return prob > threshold, prob


# ---------------------- 5. 执行单元模拟 ----------------------
class AlarmActuator:
    """声光报警执行单元"""
    def __init__(self, alarm_type='both'):
        """
        alarm_type: 'sound' (仅声音), 'light' (仅灯光), 'both' (声光)
        """
        self.alarm_type = alarm_type
        self.is_alarming = False
        
    def activate(self, alarm_level=1):
        """激活报警系统"""
        self.is_alarming = True
        alarm_thread = threading.Thread(target=self._alarm_loop, args=(alarm_level,))
        alarm_thread.daemon = True
        alarm_thread.start()
        logger.warning(f"报警系统已激活 - 级别: {alarm_level}")
        
    def deactivate(self):
        """关闭报警系统"""
        self.is_alarming = False
        logger.info("报警系统已关闭")
        
    def _alarm_loop(self, alarm_level):
        """报警循环"""
        while self.is_alarming:
            if self.alarm_type in ['sound', 'both']:
                self._trigger_sound(alarm_level)
                
            if self.alarm_type in ['light', 'both']:
                self._trigger_light(alarm_level)
                
            time.sleep(0.5)  # 报警频率
            
    def _trigger_sound(self, level):
        """触发声音报警"""
        # 实际应用中可通过扬声器播放报警音
        print(f"\033[91m{'WARNING! ' * level}\033[0m")
        
    def _trigger_light(self, level):
        """触发灯光报警"""
        # 实际应用中可通过GPIO控制LED灯
        print(f"{'💡' * level}")


# ---------------------- 6. 主系统 ----------------------
class AnomalyDetectionSystem:
    """异常行为检测系统"""
    def __init__(self, 
                 video_source=0,
                 yolo_model='yolov5s.pt',
                 action_model='i3d_r50.pth',
                 fire_model='fire_smoke_model.pt',
                 output_dir='results',
                 device='cuda'):
        
        # 初始化组件
        self.sensor = CameraSensor(video_source)
        self.detector = YOLOv5Detector(yolo_model, device)
        self.action_recognizer = ActionRecognizer(action_model, device)
        self.fire_detector = FireSmokeDetector(fire_model, device)
        self.actuator = AlarmActuator(alarm_type='both')
        
        # 结果保存
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)
        
        # 状态变量
        self.running = False
        self.tracked_persons = {}  # {track_id: [frame_history]}
        self.anomaly_events = []
        
        # 异常类型映射
        self.anomaly_types = {
            0: "正常",
            1: "打架",
            2: "摔倒",
            3: "奔跑",
            4: "烟火"
        }
        
        # 颜色映射
        self.colors = {
            0: (0, 255, 0),    # 绿色: 正常
            1: (0, 0, 255),    # 红色: 打架
            2: (0, 255, 255),  # 黄色: 摔倒
            3: (255, 0, 0),    # 蓝色: 奔跑
            4: (255, 165, 0)   # 橙色: 烟火
        }
        
    def run(self):
        """运行检测系统"""
        try:
            self.sensor.open()
            self.running = True
            self._process_frames()
        except Exception as e:
            logger.error(f"系统运行出错: {str(e)}")
        finally:
            self.sensor.release()
            self.actuator.deactivate()
            self._save_results()
            
    def _process_frames(self):
        """处理视频帧"""
        frame_count = 0
        alarm_triggered = False
        
        while self.running:
            frame = self.sensor.read_frame()
            if frame is None:
                break
                
            # 目标检测
            boxes, scores, class_ids, class_names = self.detector.detect(frame)
            
            # 过滤人员
            person_boxes = []
            for i, cls_name in enumerate(class_names):
                if cls_name == 'person':
                    person_boxes.append(boxes[i])
            
            # 行为识别
            frame_anomalies = []
            for box in person_boxes:
                x1, y1, x2, y2 = box
                person_frame = frame[y1:y2, x1:x2]
                
                if person_frame.size > 0:
                    # 更新人员跟踪历史
                    person_id = hash(f"{x1}_{y1}_{x2}_{y2}")  # 简单的ID生成
                    if person_id not in self.tracked_persons:
                        self.tracked_persons[person_id] = deque(maxlen=50)
                        
                    self.tracked_persons[person_id].append(person_frame)
                    
                    # 行为识别
                    if len(self.tracked_persons[person_id]) >= 16:
                        action_id = self.action_recognizer.recognize(
                            self.tracked_persons[person_id]
                        )
                        
                        if action_id != 0:  # 异常行为
                            frame_anomalies.append({
                                'type': action_id,
                                'box': box,
                                'person_id': person_id
                            })
                            
                            # 绘制边界框和标签
                            cv2.rectangle(frame, (x1, y1), (x2, y2), self.colors[action_id], 2)
                            cv2.putText(frame, f"{self.anomaly_types[action_id]}", 
                                       (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, 
                                       self.colors[action_id], 2)
            
            # 烟火检测
            is_fire, fire_prob = self.fire_detector.detect(frame)
            if is_fire:
                frame_anomalies.append({
                    'type': 4,
                    'box': (0, 0, frame.shape[1], frame.shape[0]),  # 整帧
                    'prob': fire_prob
                })
                cv2.putText(frame, "警告: 烟火检测", (10, 30), 
                           cv2.FONT_HERSHEY_SIMPLEX, 1, self.colors[4], 2)
            
            # 更新异常事件记录
            if frame_anomalies:
                self._record_anomaly(frame, frame_anomalies)
                
                # 触发报警
                if not alarm_triggered:
                    self.actuator.activate(alarm_level=len(frame_anomalies))
                    alarm_triggered = True
            else:
                # 无异常，关闭报警
                if alarm_triggered:
                    self.actuator.deactivate()
                    alarm_triggered = False
            
            # 显示处理后的帧
            cv2.imshow('异常行为检测系统', frame)
            
            # 按 'q' 键退出
            if cv2.waitKey(1) & 0xFF == ord('q'):
                self.running = False
                
            frame_count += 1
            if frame_count % 100 == 0:
                logger.info(f"已处理 {frame_count} 帧")
        
        cv2.destroyAllWindows()
        
    def _record_anomaly(self, frame, anomalies):
        """记录异常事件"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        event = {
            'timestamp': timestamp,
            'anomalies': [{
                'type': anomaly['type'],
                'type_name': self.anomaly_types[anomaly['type']],
                'box': anomaly.get('box', [0, 0, 0, 0]),
                'person_id': anomaly.get('person_id', None),
                'prob': anomaly.get('prob', 1.0)
            } for anomaly in anomalies]
        }
        
        self.anomaly_events.append(event)
        logger.warning(f"检测到异常: {event}")
        
        # 保存异常帧
        cv2.imwrite(f"{self.output_dir}/anomaly_{timestamp}.jpg", frame)
        
    def _save_results(self):
        """保存检测结果"""
        import json
        with open(f"{self.output_dir}/results.json", 'w') as f:
            json.dump(self.anomaly_events, f, indent=4)
            
        logger.info(f"检测结果已保存至 {self.output_dir}/results.json")


# ---------------------- 7. 主函数 ----------------------
if __name__ == "__main__":
    # 创建并运行系统
    system = AnomalyDetectionSystem(
        video_source='input_video.mp4',  # 0表示默认摄像头，也可以指定视频文件路径
        yolo_model='yolov5s.pt',
        action_model='i3d_r50.pth',
        fire_model='fire_smoke_model.pt',
        output_dir='results',
        device='cuda' if torch.cuda.is_available() else 'cpu'
    )
    
    system.run()